'''Python FFT Demo'''

import numpy as np
from typing import Tuple, Dict, List, Optional
import matplotlib.pyplot as plt
import os
import csv

class FFTDemoBase(object):
    def fft_plot(self, time_seq: np.ndarray, signal: np.ndarray, sample_rate: int, x_limit_min: Optional[float]=None, x_limit_max: Optional[float]=None, phase_only: bool=False) -> None:
        '''绘制FFT图'''
        duration = time_seq[-1] - time_seq[0] # 信号持续时间
        print(duration)

        # 对信号进行傅里叶变换
        transformed = np.fft.rfft(signal)
        # 计算绘图参数
        n = len(time_seq) # 采样数量，信号的长度
        freq = np.fft.rfftfreq(n, 1 / sample_rate) # 用户获取绘图横坐标值，传入参数分别是采样数量和采样的周期
        amplitude = np.abs(transformed) * 2 / n # 对傅里叶变换的结果取绝对值，然后进行归一化(*2/n)，得到不同频率的振幅大小

        # 相位解缠绕
        phase = np.unwrap(np.angle(transformed))
        # 相位延迟秒数计算
        phase_delay = np.array([0 if abs(freq[i]) < 10**-6 else -phase[i] / (2 * np.pi * freq[i]) for i in range(len(freq))])

        plt.figure(figsize=(12, 4))
        # 时域子图
        plt.subplot(1, 3, 1)
        plt.plot(time_seq, signal, color='blue')
        plt.title("Time Domain Signal")
        plt.xlabel("Time (s)")
        plt.ylabel("Amplitude")
        plt.xlim(0, duration)
        plt.grid(True)

        # 频域子图。
        plt.subplot(1, 3, 2)
        if x_limit_min is not None:
            plt.plot(freq, amplitude, color='red')
            plt.xlim(x_limit_min, x_limit_max)
        else:
            # 横轴对数显示：
            plt.semilogx(freq, amplitude, color='red')
        plt.title("Frequency Domain Spectrum")
        plt.xlabel("Frequency (Hz)")
        plt.ylabel("Amplitude")
        plt.grid(True)

        # 延迟子图
        plt.subplot(1, 3, 3)
        if x_limit_min is not None:
            plt.plot(freq, phase if phase_only else phase_delay, color='green')
            plt.xlim(x_limit_min, x_limit_max)
        else:
            # 横轴对数显示：
            plt.semilogx(freq, phase if phase_only else phase_delay, color='green')
        plt.title('Phase Delay')
        plt.xlabel('Frequency (Hz)')
        plt.ylabel('Phase Delay (angle)' if phase_only else 'Phase Delay (s)')
        plt.grid(True)

        plt.show()

# 傅里叶变换调试1，会先生成一个正弦信号，然后对它进行傅里叶变换并显示变换后的图
class FFTDemo1(FFTDemoBase):
    def __init__(self):
        frequency: float=10 # 生成的信号频率
        amplitude: float=1.0 # 生成的信号振幅
        sample_rate: int=100 # 采样频率
        duration: float=1.0 # 信号持续时间

        # time_seq 时间序列，按顺序均匀排列递增的时间点，单位是秒
        # sinx 正弦波序列，单位是振幅
        time_seq, sinx = self.generate_sinx(frequency, amplitude, sample_rate, duration)
        self.fft_plot(time_seq, sinx, sample_rate)
    def generate_sinx(self, frequency: float, amplitude: float, sample_rate: int, duration: float) -> Tuple[np.ndarray, np.ndarray]:
        '''生成正弦波
        函数入参：
        frequency: 频率
        amplitude: 振幅（峰峰值）
        sample_rate: 采样频率
        duration: 信号持续时间（秒）
        返回元组，其中第一个是时间序列，第二个是正弦波幅度序列
        '''
        time_seq: np.ndarray = np.linspace(0, duration, int(duration * sample_rate))
        x: np.ndarray = 2 * np.pi * frequency * time_seq
        sinx: np.ndarray = amplitude * np.sin(x)
        return time_seq, sinx

# 读取示波器采样得到的冲击响应数据文件，然后绘制出时域和频域的波形图，还有相位延迟。
class FFTDemo2(FFTDemoBase):
    def __init__(self, data_file: str):
        '''
        读取文件并绘制图表
        '''
        self.time_head = 'Time(us)'
        self.u_head = 'U(mV)'
        # 读取时间序列和电压值序列
        time_seq, u_seq = self.read_data(data_file)
        # 计算采样率
        sample_rate = len(time_seq) / (time_seq[1] - time_seq[0])
        self.fft_plot(time_seq, u_seq, sample_rate)
    def read_data(self, file_path: str) -> Tuple[np.ndarray, np.ndarray]:
        if not os.path.exists(file_path):
            raise FileNotFoundError(f'文件 {file_path} 不存在')
        headers = None
        time_s = []
        u_v = []
        with open(file_path, 'r', encoding='utf8') as f:
            reader = csv.reader(f)
            for row in reader:
                row = [item for item in row]
                if headers is None:
                    headers = row
                    # 校验headers中必须存在时间和电压两个字段
                    if self.time_head not in headers or self.u_head not in headers:
                        raise ValueError(f'文件 {file_path} 中必须存在 {self.time_head} 和 {self.u_v_head} 两个字段')
                    continue
                if len(row) < len(headers):
                    raise ValueError(f'数据行的列数少于标题行')
                data_row = {headers[i]: row[i] for i in range(len(headers))}
                time_s.append(float(data_row[self.time_head]) / 10**6)
                u_v.append(float(data_row[self.u_head]) / 10**3)
        if headers is None:
            raise ValueError(f'数据文件没有标题行')
        if len(time_s) != len(u_v):
            raise ValueError(f'时间数据和电压数据的行数不一致')
        if len(time_s) < 2:
            raise ValueError(f'数据行数少于2行') # 后续需要计算时间差，不可以少于2行
        return np.array(time_s), np.array(u_v)

# 将频率限制为20Hz-20kHz的图，相位延迟图以相位角度显示
class FFTDemo3(FFTDemo2):
    def fft_plot(self, time_seq: np.ndarray, sinx: np.ndarray, sample_rate: int, x_limit_min: Optional[float]=20, x_limit_max: Optional[float]=20000, phase_only: bool=True) -> None:
        super().fft_plot(time_seq, sinx, sample_rate, x_limit_min, x_limit_max, phase_only)

if __name__ == '__main__':
    #exit()
    #FFTDemo1()
    #FFTDemo3('data.csv')
    FFTDemo3('1_data.csv')
